import faiss
import os
import json
import logging
import numpy as np
from .embeddings import get_embedding, hidden_size

KBS_PATH = '/Users/caozy/Documents/projects/chat_driven_business/backend/kbs'
DOC_PATH =  os.path.join(KBS_PATH, "docs")
DB_PATH = os.path.join(KBS_PATH, 'index_db')
INDEX_MAPPING_PATH = os.path.join(KBS_PATH, 'index_mapping.json')

class FileNameHandler:
    def __init__(self, doc_path, surfix_descriptor=".descriptor", surfix_api=".api"):
        self.doc_path = doc_path
        self.surfix_descriptor = surfix_descriptor
        self.surfix_api = surfix_api

    def get_descriptor_filename(self, filename):
        return os.path.join(self.doc_path, filename)

    def get_api_filename(self, descriptor_name):
        if descriptor_name.startswith(self.doc_path):
            return descriptor_name.replace(self.surfix_descriptor, self.surfix_api)
        return os.path.join(self.doc_path, descriptor_name.replace(self.surfix_descriptor, self.surfix_api))

    def is_descriptor_file(self, filename):
        return filename.endswith(self.surfix_descriptor)

    def list_descriptor_names(self):
        for filename in os.listdir(self.doc_path):
            if self.is_descriptor_file(filename):
                yield self.get_descriptor_filename(filename)


kb_file_handler = FileNameHandler(DOC_PATH)

# 配置日志
logging.basicConfig(filename='build_index.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

SURFIX_DESCRIPTOR = ".descriptor"
SURFIX_API = ".api"
def load_descriptors():
    for filename in kb_file_handler.list_descriptor_names():
        with open(filename, 'r') as file:
            content = file.read()
            logging.debug("Loaded descriptor: %s", filename)
            yield (filename, content)

def build_index(descriptors):
    logging.debug("Building index")
    index = faiss.IndexFlatL2(hidden_size)
    index_mapping = {}
    
    for idx, (filename, content) in enumerate(descriptors):
        embedding = get_embedding(content)
        embedding = np.array(embedding).reshape(1, -1)
        index.add(embedding)
        index_mapping[idx] = filename
        logging.debug("Added embedding for file: %s", filename)
    
    faiss.write_index(index, DB_PATH)
    logging.debug("Index written to %s", DB_PATH)
    with open(INDEX_MAPPING_PATH, 'w') as f:
        json.dump(index_mapping, f)
    logging.debug("Index mapping written to %s", INDEX_MAPPING_PATH)

def load_doc_name(index) -> str:
    if os.path.exists(INDEX_MAPPING_PATH):
        with open(INDEX_MAPPING_PATH, 'r') as f:
            index_mapping = json.load(f)
        return index_mapping.get(str(index))
    else:
        raise FileNotFoundError(f"文件 {INDEX_MAPPING_PATH} 不存在")

def load_doc(doc_name) -> str:
    doc_file = os.path.join(DOC_PATH, doc_name)
    if os.path.exists(doc_file):
        with open(doc_file, 'r') as f:
            return f.read()
    else:
        raise FileNotFoundError(f"文件 {doc_file} 不存在")

def load_doc_content(index) -> str:
    return load_doc(load_doc_name(index))

def load_doc_api(descriptor_name):
    api_file_path = kb_file_handler.get_api_filename(descriptor_name)
    if os.path.exists(api_file_path):
        with open(api_file_path, 'r') as f:
            return f.read()
    else:
        raise FileNotFoundError(f"文件 {api_file_path} 不存在")


def load_index_db():
    """从磁盘加载向量索引"""
    if os.path.exists(DB_PATH):
        return faiss.read_index(DB_PATH)
    else:
        raise FileNotFoundError(f"文件 {DB_PATH} 不存在")


def main():
    logging.debug("Starting main function")
    descriptors = load_descriptors()
    build_index(descriptors)
    logging.debug("Main function completed")

if __name__ == '__main__':
    main()
